import torch 
import torchvision
import torch.nn.functional as F
import torchvision.transforms as transforms
import torch.optim as optim
import matplotlib.pyplot as plt
import numpy as np

from torch.utils.tensorboard import SummaryWriter

myTransforms = transforms.Compose(
    [transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
train_dataset = torchvision.datasets.CIFAR10(root='./dataset', train=True, download=True, transform=myTransforms)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=2)

test_dataset = torchvision.datasets.CIFAR10(root='./dataset', train=False, download=True, transform=myTransforms)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=4, shuffle=True, num_workers=2)


plt.imshow(train_dataset.data[2]) 
plt.show() 

'''dataiter = iter(train_loader)
images, _ = dataiter.next()
images_comb = torchvision.utils.make_grid(images)
images_comb_unnor = (images_comb*0.5+0.5).numpy()
plt.imshow(np.transpose(images_comb_unnor, (1, 2, 0)))
plt.show() '''

''' tensorboard '''
# myWriter = SummaryWriter('./log/')

class myNetwork(torch.nn.Module):
    def __init__(self):
        super(myNetwork, self).__init__()
        self.conv1 = torch.nn.Conv2d(in_channels=3, out_channels = 32, kernel_size=5, stride=1, padding=2)
        self.pool1 = torch.nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = torch.nn.Conv2d(in_channels=32, out_channels = 32, kernel_size=5, stride=1, padding=2)
        self.pool2 = torch.nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv3 = torch.nn.Conv2d(in_channels=32, out_channels = 32, kernel_size=5, stride=1, padding=2)
        self.pool3 = torch.nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc1 = torch.nn.Linear(in_features = 32*4*4, out_features = 64, bias=True)
        self.fc2 = torch.nn.Linear(in_features = 64, out_features = 64, bias=True)
        self.fc3 = torch.nn.Linear(in_features = 64, out_features = 10, bias=True)

    def forward(self, train_data):
        output = self.pool1(torch.nn.functional.relu(self.conv1(train_data)))
        output = self.pool2(torch.nn.functional.relu(self.conv2(output)))
        output = self.pool3(torch.nn.functional.relu(self.conv3(output)))
        output = output.view(-1, 32*4*4)
        output = torch.nn.functional.relu(self.fc1(output))
        output = torch.nn.functional.relu(self.fc2(output))
        output = self.fc3(output)

        return output

    
myModel = myNetwork()
myOptimzier = optim.SGD(myModel.parameters(), lr = 0.01, momentum=0.9)
myLoss = torch.nn.CrossEntropyLoss()

for _epoch in range(10):
    training_loss = 0.0
    for _step, input_data in enumerate(train_loader):

        image, label = input_data
        predict_label = myModel.forward(image)
        loss = myLoss(predict_label, label)

        # myWriter.add_scalar('training loss', loss, global_step = _epoch*len(train_loader) + _step)

        myOptimzier.zero_grad()
        loss.backward()
        myOptimzier.step()

        training_loss = training_loss + loss.item()
        if _step % 10 == 0 :
            print('[iteration - %3d] training loss: %.3f' % (_epoch*len(train_loader) + _step, training_loss/10))
            training_loss = 0.0


correct = 0
total = 0
with torch.no_grad():
    for images,labels in test_loader:
        outputs = myModel(images)
        numbers,predicted = torch.max(outputs.data,1)
        total += labels.size(0)
        correct += (predicted==labels).sum().item()

print('Testing Accuracy : %d %%' % ( 100 * correct / total))




